import socket
import os
import json
import threading
from shared.crypto import CryptoUtils
from shared.database import DatabaseManager
import stat

class FTPServer:
    """
    FTP服务器类
    实现了一个支持加密通信的FTP服务器，包含以下主要功能：
    - 用户认证和管理
    - 文件上传和下载
    - 目录操作（创建、切换、列表等）
    - 实时聊天功能
    - 在线用户管理
    """
    def __init__(self, host='localhost', port=21000):
        # 初始化服务器socket配置
        self.host = host
        self.port = port
        self.server_socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        self.server_socket.bind((self.host, self.port))
        self.server_socket.listen(5)

        # 生成RSA公钥和私钥（服务器启动时执行一次）
        self.private_key, self.public_key = CryptoUtils.generate_key_pair()
        # 初始化数据库管理器
        self.db = DatabaseManager()
        # 设置FTP根目录
        self.root_dir = os.path.join(os.path.dirname(__file__), 'ftp_root')
        os.makedirs(self.root_dir, exist_ok=True)
        
        # 启动时往数据库添加默认管理员账户
        try:
            self.db.add_user('admin', 'admin123')
            print("已创建或已存在默认admin用户")
        except Exception as e:
            print(f"错误: {str(e)}")
            
        self.clients = {}  # 存储已连接的客户端信息 {username: (socket, symmetric_key)}
        print(f"服务器启动在 {host}:{port}")

    
    def handle_client(self, client_socket, client_address):
        """
        处理单个客户端连接的主方法
        实现了客户端的认证、通信加密和命令处理流程

        Args:
            client_socket (socket): 客户端socket连接
            client_address (tuple): 客户端地址信息(ip, port)
        """
        username = None
        symmetric_key = None
        try:
            # 发送公钥的PEM格式给客户端
            client_socket.send(CryptoUtils.serialize_public_key(self.public_key))

            # 接收对称密钥进行解密
            encrypted_symmetric_key = client_socket.recv(2048)
            symmetric_key = CryptoUtils.decrypt_with_private_key(
                self.private_key, encrypted_symmetric_key)

            # 处理认证
            auth_result, username = self._handle_authentication(client_socket, symmetric_key)
            if not auth_result:
                return
                
            self.clients[username] = (client_socket, symmetric_key)

            current_dir = self.root_dir
            while True:
                encrypted_data = client_socket.recv(4096)
                if not encrypted_data:
                    break

                decrypted_data = CryptoUtils.decrypt_message(symmetric_key, encrypted_data)
                command_data = json.loads(decrypted_data)
                command = command_data['command']
                args = command_data.get('args', [])

                response = self._handle_command(command, args, current_dir, username, client_socket, symmetric_key)
                encrypted_response = CryptoUtils.encrypt_message(
                    symmetric_key, json.dumps(response))
                client_socket.send(encrypted_response)

                if command == 'cd' and response['success']:
                    current_dir = response['data']

        except Exception as e:
            print(f"Error handling client {client_address}: {str(e)}")
        finally:
            if username and username in self.clients:
                del self.clients[username]
            client_socket.close()

    def _handle_authentication(self, client_socket, symmetric_key):
        """
        处理客户端的身份认证

        Args:
            client_socket (socket): 客户端socket连接
            symmetric_key (bytes): 用于加密通信的对称密钥

        Returns:
            tuple: (认证结果(bool), 用户名(str))
        """
        encrypted_credentials = client_socket.recv(1024)
        decrypted_credentials = CryptoUtils.decrypt_message(
            symmetric_key, encrypted_credentials)
        credentials = json.loads(decrypted_credentials)

        auth_result = self.db.verify_user(
            credentials['username'], credentials['password'])
        
        response = {'success': auth_result}
        encrypted_response = CryptoUtils.encrypt_message(
            symmetric_key, json.dumps(response))
        client_socket.send(encrypted_response)
        return auth_result, credentials['username'] if auth_result else None

    def _handle_command(self, command, args, current_dir, username=None, client_socket=None, symmetric_key=None):
        """
        处理客户端发送的各种命令
        支持文件操作、目录操作、聊天等功能

        Args:
            command (str): 命令名称
            args (list): 命令参数列表
            current_dir (str): 当前工作目录
            username (str, optional): 用户名
            client_socket (socket, optional): 客户端socket连接
            symmetric_key (bytes, optional): 加密密钥

        Returns:
            dict: 命令执行结果
        """
        handlers = {
            'list': self._handle_list,
            'cd': self._handle_cd,
            'pwd': self._handle_pwd,
            'mkdir': self._handle_mkdir,
            'get': self._handle_get,
            'put': self._handle_put,
            'attrib': self._handle_attrib,
            'rename': self._handle_rename,
            'create': self._handle_create,
            'add': self._handle_add,
            'online': self._handle_online
        }

        # 处理聊天命令
        if command == 'chat' and username and client_socket and symmetric_key:
            return self._handle_chat(args, current_dir, username, client_socket, symmetric_key)

        # 处理在线用户命令
        if command == 'online':
            response = self._handle_online(args, current_dir)
            response['type'] = 'command_response'
            return response

        # 处理其他命令
        handler = handlers.get(command)
        if handler:
            response = handler(args, current_dir)
            response['type'] = 'command_response'
            return response
        return {'success': False, 'message': 'Unknown command', 'type': 'command_response'}

    def _handle_list(self, args, current_dir):
        """
        处理ls/dir命令，列出当前目录下的文件和文件夹

        Args:
            args (list): 命令参数（此方法不使用参数）
            current_dir (str): 当前工作目录

        Returns:
            dict: 包含目录内容列表的响应字典
        """
        try:
            items = os.listdir(current_dir)
            return {'success': True, 'data': items}
        except Exception as e:
            return {'success': False, 'message': str(e)}

    def _handle_cd(self, args, current_dir):
        """
        处理cd命令，切换当前工作目录

        Args:
            args (list): 目标目录路径
            current_dir (str): 当前工作目录

        Returns:
            dict: 包含新目录路径的响应字典
        """
        try:
            if not args:
                new_path = self.root_dir
            elif args[0] == '..':
                new_path = os.path.dirname(current_dir)
                if not new_path.startswith(self.root_dir):
                    new_path = self.root_dir
            else:
                new_path = os.path.join(current_dir, args[0])

            if os.path.isdir(new_path) and new_path.startswith(self.root_dir):
                return {'success': True, 'data': new_path}
            return {'success': False, 'message': 'Invalid directory'}
        except Exception as e:
            return {'success': False, 'message': str(e)}

    def _handle_pwd(self, args, current_dir):
        """
        处理pwd命令，显示当前工作目录

        Args:
            args (list): 命令参数（此方法不使用参数）
            current_dir (str): 当前工作目录

        Returns:
            dict: 包含当前目录相对路径的响应字典
        """
        rel_path = os.path.relpath(current_dir, self.root_dir)
        return {'success': True, 'data': rel_path}

    def _handle_mkdir(self, args, current_dir):
        """
        处理mkdir命令，创建新目录

        Args:
            args (list): 新目录名
            current_dir (str): 当前工作目录

        Returns:
            dict: 操作结果响应字典
        """
        if not args:
            return {'success': False, 'message': 'Directory name required'}
        try:
            new_dir = os.path.join(current_dir, args[0])
            os.makedirs(new_dir)
            return {'success': True}
        except Exception as e:
            return {'success': False, 'message': str(e)}

    def _handle_get(self, args, current_dir):
        """
        处理get命令，实现文件下载功能

        Args:
            args (list): 要下载的文件名
            current_dir (str): 当前工作目录

        Returns:
            dict: 包含文件内容（十六进制格式）的响应字典
        """
        if not args:
            return {'success': False, 'message': 'Filename required'}
        try:
            file_path = os.path.join(current_dir, args[0])
            if not os.path.isfile(file_path):
                return {'success': False, 'message': 'File not found'}
            with open(file_path, 'rb') as f:
                data = f.read()
            return {'success': True, 'data': data.hex()}
        except Exception as e:
            return {'success': False, 'message': str(e)}

    def _handle_put(self, args, current_dir):
        """
        处理put命令，实现文件上传功能

        Args:
            args (list): [文件名, 文件内容（十六进制格式）]
            current_dir (str): 当前工作目录

        Returns:
            dict: 操作结果响应字典
        """
        if len(args) != 2:
            return {'success': False, 'message': 'Filename and data required'}
        try:
            file_path = os.path.join(current_dir, args[0])
            data = bytes.fromhex(args[1])
            with open(file_path, 'wb') as f:
                f.write(data)
            return {'success': True}
        except Exception as e:
            return {'success': False, 'message': str(e)}

    def _handle_attrib(self, args, current_dir):
        """
        处理attrib命令，获取文件属性信息

        Args:
            args (list): 目标文件名
            current_dir (str): 当前工作目录

        Returns:
            dict: 包含文件属性信息的响应字典
        """
        if not args:
            return {'success': False, 'message': 'Filename required'}
        try:
            file_path = os.path.join(current_dir, args[0])
            if not os.path.exists(file_path):
                return {'success': False, 'message': 'File not found'}
            
            stats = os.stat(file_path)
            attrs = {
                'size': stats.st_size,
                'created': stats.st_ctime,
                'modified': stats.st_mtime,
                'mode': stat.filemode(stats.st_mode)
            }
            return {'success': True, 'data': attrs}
        except Exception as e:
            return {'success': False, 'message': str(e)}

    def _handle_rename(self, args, current_dir):
        """
        处理rename命令，重命名文件或目录

        Args:
            args (list): [原名称, 新名称]
            current_dir (str): 当前工作目录

        Returns:
            dict: 操作结果响应字典
        """
        if len(args) != 2:
            return {'success': False, 'message': 'Old and new names required'}
        try:
            old_path = os.path.join(current_dir, args[0])
            new_path = os.path.join(current_dir, args[1])
            os.rename(old_path, new_path)
            return {'success': True}
        except Exception as e:
            return {'success': False, 'message': str(e)}

    def _handle_create(self, args, current_dir):
        """
        处理create命令，创建空文件

        Args:
            args (list): 新文件名
            current_dir (str): 当前工作目录

        Returns:
            dict: 操作结果响应字典
        """
        if not args:
            return {'success': False, 'message': 'Filename required'}
        try:
            file_path = os.path.join(current_dir, args[0])
            with open(file_path, 'w') as f:
                pass
            return {'success': True}
        except Exception as e:
            return {'success': False, 'message': str(e)}

    def _handle_add(self, args, current_dir):
        """
        处理add命令，添加新用户

        Args:
            args (list): [用户名, 密码]
            current_dir (str): 当前工作目录

        Returns:
            dict: 用户添加结果响应字典
        """
        if len(args) != 2:
            return {'success': False, 'message': 'Username and password required'}
        try:
            success = self.db.add_user(args[0], args[1])
            return {'success': success, 'message': 'User added successfully' if success else 'Failed to add user'}
        except Exception as e:
            return {'success': False, 'message': str(e)}



    def _handle_chat(self, args, current_dir, username, client_socket, symmetric_key):
        """
        处理客户端发送的聊天命令
        支持私聊和群发消息功能

        Args:
            args (list): 命令参数列表，第一个参数为接收者，后续参数组成消息内容
            current_dir (str): 当前工作目录
            username (str): 发送消息的用户名
            client_socket (socket): 客户端socket连接
            symmetric_key (bytes): 用于加密通信的对称密钥

        Returns:
            dict: 包含操作结果的字典，success字段表示是否成功，message字段包含响应信息
        """
        if len(args) < 2:
            return {'success': False, 'message': 'Please specify recipient and message'}
        
        recipient = args[0]
        message = ' '.join(args[1:])
        
        if recipient == 'all':
            # 广播消息给所有在线用户
            for user, (sock, key) in self.clients.items():
                if user != username:  # 不发给自己
                    try:
                        chat_data = {
                            'type': 'chat',
                            'from': username,
                            'message': message
                        }
                        encrypted_chat = CryptoUtils.encrypt_message(
                            key, json.dumps(chat_data))
                        sock.send(encrypted_chat)
                    except Exception as e:
                        print(f"Error sending message to {user}: {str(e)}")
            return {'success': True, 'message': 'Message broadcast successfully'}
        
        elif recipient in self.clients:
            # 发送私人消息
            recipient_socket, recipient_key = self.clients[recipient]
            try:
                chat_data = {
                    'type': 'chat',
                    'from': username,
                    'message': message
                }
                encrypted_chat = CryptoUtils.encrypt_message(
                    recipient_key, json.dumps(chat_data))
                recipient_socket.send(encrypted_chat)
                return {'success': True, 'message': f'Message sent to {recipient}'}
            except Exception as e:
                return {'success': False, 'message': f'Failed to send message: {str(e)}'}
        else:
            return {'success': False, 'message': f'User {recipient} is not online'}

    def _handle_online(self, args, current_dir):
        """
        获取当前在线用户列表

        Args:
            args (list): 命令参数列表（此方法不使用参数）
            current_dir (str): 当前工作目录

        Returns:
            dict: 包含在线用户列表的响应字典
        """
        # 获取在线用户列表
        online_users = list(self.clients.keys())
        return {'success': True, 'data': online_users}



    def run(self):
        """
        启动FTP服务器
        - 持续监听新的客户端连接
        - 为每个客户端创建新的处理线程
        - 支持通过Ctrl+C优雅关闭服务器
        """
        try:
            while True:
                client_socket, client_address = self.server_socket.accept()
                print(f"New connection from {client_address}")
                client_thread = threading.Thread(
                    target=self.handle_client,
                    args=(client_socket, client_address)
                )
                client_thread.start()
        except KeyboardInterrupt:
            print("\nServer shutting down...")
        finally:
            self.server_socket.close()
            self.db.close()

if __name__ == "__main__":
    server = FTPServer()
    server.run()
